package com.example.demo.filter;

import com.example.demo.common.AESEncryptor;
import com.example.demo.common.constant.ResultCode;
import com.example.demo.exception.ValidException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.HandlerExceptionResolver;

import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;

@WebFilter(urlPatterns = {"/password/getData"})
@Slf4j
@RequiredArgsConstructor
public class AesDecryptFilter implements Filter {
    private final HandlerExceptionResolver handlerExceptionResolver;

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain chain) throws ServletException, IOException {
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;
        String contentType = request.getContentType();

        if (contentType != null) {
            try {
                if (contentType.equals("application/json")) {
                    log.info("start filter of the AES decrypt: DTO");

                    // 获取参数（DTO型），并解密
                    String body = AESEncryptor.decrypt(request.getReader().lines().collect(Collectors.joining(System.lineSeparator())));

                    AesDecryptFilter.ModifyBodyWrapper bodyWrapper = new AesDecryptFilter.ModifyBodyWrapper(request, body);
                    chain.doFilter(bodyWrapper, servletResponse);
                }
                else if (contentType.equals("application/x-www-form-urlencoded")) {
                    log.info("start filter of the AES decrypt: Request");

                    // 获取参数（Request型）
                    AesDecryptFilter.ModifyRequestWrapper requestWrapper = new AesDecryptFilter.ModifyRequestWrapper(request);
                    Map<String, String[]> paramNameList = requestWrapper.getParameterMap();
                    for (String paramName : paramNameList.keySet()) {
                        // 将 paramValue 解密
                        String paramValue = requestWrapper.getParameterValues(paramName)[0];
                        requestWrapper.getParameterMap().put(paramName, new String[]{AESEncryptor.decrypt(paramValue)});
                    }

                    // 修改参数（只能修改原有参数的值，不能新增参数）
                    //requestWrapper.getParameterMap().put("参数名", new String[]{"参数值"});

                    chain.doFilter(requestWrapper, servletResponse);
                }
            }
            catch (Exception ex) {
                handlerExceptionResolver.resolveException(request, response, null, new ValidException(ResultCode.PARAM_RSA_ERROR));
            }
        }
        else {
            chain.doFilter(servletRequest, servletResponse);
        }
    }

    private static class ModifyBodyWrapper extends HttpServletRequestWrapper {
        private final String body;

        public ModifyBodyWrapper(HttpServletRequest request, String oldBody) {
            super(request);
            body = oldBody;
        }

        @Override
        public BufferedReader getReader() throws IOException {
            return new BufferedReader(new InputStreamReader(getInputStream()));
        }

        @Override
        public ServletInputStream getInputStream() throws IOException {
            final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(body.getBytes());
            return new ServletInputStream() {
                @Override
                public boolean isFinished() {
                    return false;
                }

                @Override
                public boolean isReady() {
                    return false;
                }

                @Override
                public void setReadListener(ReadListener readListener) {
                }

                @Override
                public int read() {
                    return byteArrayInputStream.read();
                }
            };
        }
    }

    private static class ModifyRequestWrapper extends HttpServletRequestWrapper {
        private final Map<String, String[]> parameterMap;

        public ModifyRequestWrapper(HttpServletRequest request) {
            super(request);
            parameterMap = new HashMap<>();
            // 复制原始参数
            Enumeration<String> paramNames = request.getParameterNames();
            while (paramNames.hasMoreElements()) {
                String paramName = paramNames.nextElement();
                parameterMap.put(paramName, request.getParameterValues(paramName));
            }
        }

        @Override
        public String getParameter(String name) {
            if (parameterMap.containsKey(name)) {
                return parameterMap.get(name)[0];
            }
            return super.getParameter(name);
        }

        @Override
        public Map<String, String[]> getParameterMap() {
            return parameterMap;
        }

        @Override
        public String[] getParameterValues(String name) {
            return parameterMap.get(name);
        }
    }
}
